#include <iostream>
#include <opencv2/opencv.hpp>
#include <fstream>
#include <iomanip>
#include <filesystem>
#include <vector>
#include <string>

#include "co_detr_ros/tensorrt_detector.h"

namespace fs = std::filesystem;

int main(int argc, char** argv) {
    // Configuration paths
    std::string model_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/models/co_detr_mixed_precision.plan";
    std::string plugin_path = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/grid_sampler/build/libtrtgrid_sampler.so";
    std::string data_dir = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/data";
    std::string output_dir = "/data2/xd/Co-DETR-TensorRT/co_detr_ros_ws/src/co_detr_ros/results";
    
    // Create output directory if it doesn't exist
    fs::create_directories(output_dir);
    
    // Initialize TensorRT detector
    std::cout << "Initializing Co-DETR TensorRT Detector..." << std::endl;
    TensorRTDetector detector(model_path, plugin_path);
    if (!detector.initialize()) {
        std::cerr << "Failed to initialize TensorRT detector" << std::endl;
        return -1;
    }
    
    // Get all image files in data directory
    std::vector<std::string> image_files;
    if (fs::exists(data_dir) && fs::is_directory(data_dir)) {
        for (const auto& entry : fs::directory_iterator(data_dir)) {
            if (entry.is_regular_file()) {
                std::string extension = entry.path().extension().string();
                // Support image file formats
                if (extension == ".jpg" || extension == ".jpeg" || extension == ".png" || extension == ".bmp") {
                    image_files.push_back(entry.path().string());
                }
            }
        }
    }
    
    if (image_files.empty()) {
        std::cerr << "No image files found in data directory: " << data_dir << std::endl;
        return -1;
    }
    
    std::cout << "Found " << image_files.size() << " image files for processing" << std::endl;
    
    // Process each image file
    for (size_t i = 0; i < image_files.size(); ++i) {
        const auto& input_file = image_files[i];
        
        std::cout << "\n=== Processing image: " << input_file << " ===" << std::endl;
        
        // Perform inference on image
        std::vector<float> output_data = detector.inferenceFromImage(input_file);
        if (output_data.empty()) {
            std::cout << "Inference failed for: " << input_file << std::endl;
            continue;
        }
        
        // Generate output filename
        fs::path input_path(input_file);
        std::string base_name = input_path.stem().string();
        std::string output_image_path = output_dir + "/" + base_name + "_detection_result.jpg";
        
        // Create visualization
        bool viz_success = detector.visualizeResults(output_data, output_image_path);
        if (viz_success) {
            std::cout << "✅ Detection result saved to: " << output_image_path << std::endl;
        } else {
            std::cout << "❌ Failed to save detection result for: " << input_file << std::endl;
        }
        
        std::cout << "Completed processing: " << input_file << std::endl;
    }
    
    std::cout << "\n=== All images processed successfully ===" << std::endl;
    return 0;
}
